// Manager for update operations.

#include "UpdateManager.h"
#include <iostream>
#include <iomanip>
#include "hack.h"
#include "urlmon.h"
#include <wininet.h>
#include <sstream>
#include <future>
#include "Shlwapi.h"

// Handle any pending update.
bool UpdateManager::HandlePendingUpdate(const string& currentExePath, bool& restart)
{
	restart = false;
	if (ProcessWasUpdated())
	{
		LogManager::DisplayWithSuspense(L"Update in progress");
		const string initialExePath = GetInitialExePath(currentExePath);
		if (CopyFileW(_W(currentExePath).c_str(), _W(initialExePath).c_str(), false) == 0) return LogManager::DisplayWarning(L"Error while replacing the previous executable : 0x%x.\n", GetLastError());
		return TryRestart(initialExePath, __APPLICATION_RESTARTED__, restart);
	}
	else if (ProcessWasRestarted())
	{
		wcout << LIGHTGREEN << L"Update successful !\n" << LIGHTGREY;
		_wremove(_W(GetBackupExePath(currentExePath)).c_str());
	}

	return true;
}

// Try to update if it is mandatory.
bool UpdateManager::TryUpdateIfMandatory(ConnectionInformations& connectionInformations, VersionInformations& versionInformations, bool& restart)
{
	restart = false;

	// check if update is mandatory
	if (IsUpdateMandatory(connectionInformations.RequiredVersion, versionInformations.ProductVersion))
	{
		wcout << LIGHTRED << L"You must use the version " << YELLOW << _W(connectionInformations.RequiredVersion) << LIGHTRED << L" to connect to the server.\n The update will automatically be downloaded right now.\n" << LIGHTGREY;

		const string backupExePath = GetBackupExePath(versionInformations.CurrentExePath);

		// try to update then try to restart
		return TryDownload(connectionInformations.UpdateAddress, backupExePath) && TryRestart(backupExePath, __APPLICATION_UPDATED__, restart);
	}

	// no update
	return true;
}

// Indicates wether an update is mandatory.
bool UpdateManager::IsUpdateMandatory(const string& requiredVersion, const string& currentVersion)
{
	vector<string> vRequiredVersion = Split(requiredVersion, '.');
	vector<string> vCurrentVersion = Split(currentVersion, '.');
	const size_t size = vRequiredVersion.size();
	for (size_t i = 0; i < size; i++)
		if (atoi(vRequiredVersion[i].c_str()) > atoi(vCurrentVersion[i].c_str()))
			return true;
	return false;
}

// Try to restart the software.
bool UpdateManager::TryRestart(const string& backupExePath, const string& commandLineSwitch, bool& restart)
{
	LogManager::DisplayWithSuspense(L"Restarting");
	STARTUPINFOW startupInfo = {0};
	GetStartupInfoW(&startupInfo);
	startupInfo.cb = sizeof(STARTUPINFOW);

	PROCESS_INFORMATION	processInformation = {0};
	wstring commandLine = _W(backupExePath);
	PathQuoteSpacesW(&commandLine[0]);
	commandLine.append(_W(commandLineSwitch));
	if (CreateProcessW(NULL, &commandLine[0], NULL, NULL, FALSE, CREATE_NEW_CONSOLE, NULL, NULL, &startupInfo, &processInformation) == 0) return LogManager::DisplayWarning(L"Could start the new version (error : %d).\n", GetLastError());
	restart = true;
	return true;
}

// Split a value into individual components using the specified delimiter.
vector<string> UpdateManager::Split(const string &value, char delimiter)
{
	vector<string> result;
	stringstream stream(value);
	string item;
	while (getline(stream, item, delimiter)) result.push_back(item);
	return result;
}

// Try to download the specified file to the specified location.
bool UpdateManager::TryDownload(const string& fileUrl, const string& destinationPath)
{
	LPCWSTR url = _W(fileUrl).c_str();

	// invalidate the url cache to force download
	DeleteUrlCacheEntryW(url);

	// try to download the file
	HRESULT hr = URLDownloadToFileW(NULL, url, _W(destinationPath).c_str(), 0, NULL);
	delete[] url;

	// download success
	if(SUCCEEDED(hr)) return true;

	// download failed
	wstringstream stream;
	stream << L"Could not download the update (error code : 0x, file url : %ls).\n" << hex << hr << _W(fileUrl).c_str() << endl;

	return LogManager::DisplayWarning(stream.str().c_str());
}

// Indicates whether the process was updated.
bool UpdateManager::ProcessWasUpdated()
{
	return StrStrW(GetCommandLineW(), _W(__APPLICATION_UPDATED__).c_str()) != NULL;
}

// Indicates whether the process was restarted.
bool UpdateManager::ProcessWasRestarted()
{
	return StrStrW(GetCommandLineW(), _W(__APPLICATION_RESTARTED__).c_str()) != NULL;
}

// Gets the initial executabale path.
string UpdateManager::GetInitialExePath(const string& currentExePath)
{
	string initialExePath = currentExePath;
	return initialExePath.replace(initialExePath.size() - 8, initialExePath.size(), ".exe");
}

// Gets the backup executabale path.
string UpdateManager::GetBackupExePath(const string& currentExePath)
{
	string backupExePath = currentExePath;
	return backupExePath.replace(backupExePath.size() - 4, backupExePath.size(), ".old.exe");
}